{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Colab_DAIN.ipynb",
      "provenance": [],
      "private_outputs": true,
      "collapsed_sections": [],
      "toc_visible": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1pIo4r_Y8cMo",
        "colab_type": "text"
      },
      "source": [
        "# DAIN Colab"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iGPHW5SOpPe3",
        "colab_type": "text"
      },
      "source": [
        "Credits to original Colab file: \n",
        "https://github.com/baowenbo/DAIN/issues/44\n",
        "\n",
        "My fork:\n",
        "https://github.com/styler00dollar/DAIN\n",
        "\n",
        "Enhancement by [Styler00Dollar](https://github.com/styler00dollar) aka \"sudo rm -rf / --no-preserve-root#8353\" on discord and [Alpha](https://github.com/AlphaGit). Please do not run this command in your linux terminal. It's rather meant as a joke.\n",
        "\n",
        "A simple guide:\n",
        "- Copy the .ipynb-file to your drive.\n",
        "- Create a folder inside of Google Drive named \"DAIN\"\n",
        "- Change the configurations in the next cell\n",
        "- Run cells one by one\n",
        "\n",
        "Stuff that should be improved:\n",
        "- Alpha channel will be removed automatically and won't be added back. Anything related to alpha will be converted to black.\n",
        "- Adding configuration to select speed\n",
        "- Detect scenes to avoid interpolating scene-changes\n",
        "- Auto-resume\n",
        "- Copy `start_frame` - `end_frame` audio from original input to final output"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "enKoi0TR2fOD",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "################# Configurations ############################\n",
        "\n",
        "# Use the values in here to configure what you'd like DAIN to do.\n",
        "\n",
        "# Input file: Path (relative to the root of your Google Drive) to the input file.\n",
        "# For instance, if you save your \"example.mkv\" file in your Google Drive, inside a \"videos\" folder, the path would be:\n",
        "# videos/example.mkv. Currenly videos and gifs are supported.\n",
        "INPUT_FILEPATH = \"DAIN/input.mp4\"\n",
        "\n",
        "# Output file path: path (relative to the root of your Google Drive) for the output file. It will also determine the\n",
        "# filetype in the destination. MP4 is recommended.\n",
        "OUTPUT_FILE_PATH = \"DAIN/output.mp4\"\n",
        "\n",
        "################# Optional configurations ############################\n",
        "\n",
        "# Target FPS = how many frames per second should the result have. This will determine how many intermediate images are\n",
        "# interpolated.\n",
        "TARGET_FPS = 60\n",
        "\n",
        "# Frame input directly\n",
        "# Use a path that is in your GDrive if you already have the list of frames in the format 00001.png, 00002.png, etc.\n",
        "# Your GDrive is located at `/content/gdrive/My Drive/`\n",
        "FRAME_INPUT_DIR = '/content/DAIN/input_frames'\n",
        "\n",
        "# Frame output directory\n",
        "# Use a location in your GDrive if you want the generated frames stored to your Google Drive.\n",
        "# Your GDrive is located at `/content/gdrive/My Drive/DAIN/tmp`\n",
        "FRAME_OUTPUT_DIR = '/content/DAIN/output_frames'\n",
        "\n",
        "# Seamless playback\n",
        "# Creates a seamless loop by using the first frame as last one as well. Set this to True this if loop is intended.\n",
        "SEAMLESS = False\n",
        "\n",
        "# Resize hotfix\n",
        "# DAIN frames are a bit \"shifted / smaller\" compared to original input frames. This can partly be mitigated with resizing\n",
        "# DAIN frames to the resolution +2px and cropping the result to the original resoultion with the starting point (1,1).\n",
        "# Without this fix, DAIN tends to make \"vibrating\" output and it is pretty noticible with static elements like text.\n",
        "# This hotfix tries to make such effects less visible for a smoother video playback. I do not know what DAINAPP\n",
        "# uses as a fix for this problem, but the original does show such behaviour with the default test images. More advanced\n",
        "# users can change the interpolation method. The methods cv2.INTER_CUBIC and cv2.INTER_LANCZOS4 are recommended.\n",
        "# The current default value is cv2.INTER_LANCZOS4.\n",
        "RESIZE_HOTFIX = True\n",
        "\n",
        "# Auto-delete output PNG dir after ffmpeg video creation.\n",
        "AUTO_REMOVE = True"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "N9cGwalNeyk9",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Connect Google Drive\n",
        "from google.colab import drive\n",
        "drive.mount('/content/gdrive')\n",
        "print('Google Drive connected.')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "irzjv1x4e3S4",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Check your current GPU\n",
        "# If you are lucky, you get 16GB VRAM. If you are not lucky, you get less. VRAM is important. The more VRAM, the higher the maximum resolution will go.\n",
        "\n",
        "# 16GB: Can handle 720p. 1080p will procude an out-of-memory error. \n",
        "# 8GB: Can handle 480p. 720p will produce an out-of-memory error.\n",
        "\n",
        "# P100: 16GB (Works)\n",
        "# T4: 16GB [RuntimeError: CUDA call failed]\n",
        "# P4: 8GB (Works)\n",
        "# K80: 8GB (Not tested)\n",
        "\n",
        "!nvidia-smi --query-gpu=gpu_name,driver_version,memory.total --format=csv"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UYHTTP91oMvh",
        "colab_type": "text"
      },
      "source": [
        "# Install dependencies.\n",
        "\n",
        "This next step may take somewhere between 15-20 minutes. Run this only once at startup.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UeaU8um5-2NS",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from IPython.display import clear_output\n",
        "!git clone https://github.com/styler00dollar/Colab-DAIN /content/DAIN\n",
        "\n",
        "# This takes a while. Just wait. ~15 minutes.\n",
        "# Building DAIN.\n",
        "%cd /content/DAIN/my_package/\n",
        "!./build.sh\n",
        "print(\"Building #1 done.\")\n",
        "\n",
        "# Wait again. ~5 minutes.\n",
        "# Building DAIN PyTorch correlation package.\n",
        "%cd /content/DAIN/PWCNet/correlation_package_pytorch1_0\n",
        "!./build.sh\n",
        "print(\"Building #2 done.\")\n",
        "\n",
        "# Downloading pre-trained model\n",
        "%cd /content/DAIN\n",
        "!mkdir model_weights\n",
        "!wget -O model_weights/best.pth http://vllab1.ucmerced.edu/~wenbobao/DAIN/best.pth\n",
        "\n",
        "!CUDA_VISIBLE_DEVICES=0\n",
        "\n",
        "!sudo apt-get install imagemagick imagemagick-doc"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zm5kn6vTncL4",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Detecting FPS of input file.\n",
        "%shell yes | cp -f /content/gdrive/My\\ Drive/{INPUT_FILEPATH} /content/DAIN/\n",
        "\n",
        "import os\n",
        "filename = os.path.basename(INPUT_FILEPATH)\n",
        "\n",
        "import cv2\n",
        "cap = cv2.VideoCapture(f'/content/DAIN/{filename}')\n",
        "\n",
        "fps = cap.get(cv2.CAP_PROP_FPS)\n",
        "\n",
        "if(fps/TARGET_FPS>0.5):\n",
        "  print(\"Define a higher fps, because there is not enough time for new frames. (Old FPS)/(New FPS) should be lower than 0.5. Interpolation will fail if you try.\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9YNva-GuKq4Y",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# ffmpeg extract - Generating individual frame PNGs from the source file.\n",
        "%shell rm -rf '{FRAME_INPUT_DIR}'\n",
        "%shell mkdir -p '{FRAME_INPUT_DIR}'\n",
        "\n",
        "%shell ffmpeg -i '/content/DAIN/{filename}' '{FRAME_INPUT_DIR}/%05d.png'\n",
        "\n",
        "png_generated_count_command_result = %shell ls '{FRAME_INPUT_DIR}' | wc -l\n",
        "clear_output()\n",
        "\n",
        "pngs_generated_count = int(png_generated_count_command_result.output.strip())\n",
        "\n",
        "import shutil\n",
        "if SEAMLESS==True:\n",
        "  pngs_generated_count += 1\n",
        "  original = str(FRAME_INPUT_DIR)+\"/00001.png\"\n",
        "  target = str(FRAME_INPUT_DIR)+\"/\"+str(pngs_generated_count).zfill(5)+\".png\"\n",
        "  shutil.copyfile(original, target)\n",
        "\n",
        "print(f\"Input FPS: {fps}\")\n",
        "print(f\"{pngs_generated_count} frame PNGs generated.\")\n",
        "\n",
        "# Checking if PNG do have alpha\n",
        "import subprocess as sp\n",
        "%cd {FRAME_INPUT_DIR}\n",
        "channels = sp.getoutput('identify -format %[channels] 00001.png')\n",
        "print (f\"{channels} detected\")\n",
        "\n",
        "# Removing alpha if detected\n",
        "if \"a\" in channels:\n",
        "  print(\"Alpha detected and will be removed.\")\n",
        "  print(sp.getoutput('find . -name \"*.png\" -exec convert \"{}\" -alpha off PNG24:\"{}\" \\;'))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "W3rrE7L824gL",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Interpolation\n",
        "%shell mkdir -p '{FRAME_OUTPUT_DIR}'\n",
        "%cd /content/DAIN\n",
        "\n",
        "!python colab_interpolate.py --netName DAIN_slowmotion --time_step {fps/TARGET_FPS} --start_frame 1 --end_frame {pngs_generated_count} --frame_input_dir '{FRAME_INPUT_DIR}' --frame_output_dir '{FRAME_OUTPUT_DIR}'"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xPYewi8jsq_B",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Finding DAIN Frames, upscaling and cropping to match original\n",
        "import numpy as np\n",
        "%cd {FRAME_OUTPUT_DIR}\n",
        "if(RESIZE_HOTFIX==True):\n",
        "  images = []\n",
        "  for filename in os.listdir(f'{FRAME_OUTPUT_DIR}'):\n",
        "    img = cv2.imread(os.path.join(f'{FRAME_OUTPUT_DIR}',filename))\n",
        "    part_filename = os.path.splitext(filename)\n",
        "    if(part_filename[0].endswith('0')==False):\n",
        "      dimension = (img.shape[1]+2, img.shape[0]+2)\n",
        "      resized = cv2.resize(img, dimension, interpolation=cv2.INTER_LANCZOS4)\n",
        "      crop = resized[1:(dimension[1]-1), 1:(dimension[0]-1)]\n",
        "      cv2.imwrite(part_filename[0]+\".png\", crop)\n",
        "\n",
        "%cd /content/DAIN"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TKREDli2IDMV",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Create video\n",
        "%cd {FRAME_OUTPUT_DIR}\n",
        "%shell ffmpeg -y -r {TARGET_FPS} -f image2 -pattern_type glob -i '*.png' '/content/gdrive/My Drive/{OUTPUT_FILE_PATH}'\n",
        "if(AUTO_REMOVE==True):\n",
        "  !rm -rf {FRAME_OUTPUT_DIR}/*"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uCBHIXNN-JYu",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# [Experimental] Create video with sound\n",
        "# Only run this, if the original had sound.\n",
        "%cd {FRAME_OUTPUT_DIR}\n",
        "%shell ffmpeg -i '/content/DAIN/{filename}' -acodec copy output-audio.aac\n",
        "%shell ffmpeg -y -r {TARGET_FPS} -f image2 -pattern_type glob -i '*.png' -i output-audio.aac -shortest '/content/gdrive/My Drive/{OUTPUT_FILE_PATH}'\n",
        "if(AUTO_REMOVE==True):\n",
        "  !rm -rf {FRAME_OUTPUT_DIR}/*\n",
        "  !rm -rf output-audio.aac"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8tTXpWkRBm8T",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Manually removing output PNG dir (before working with another file)\n",
        "!rm -rf {FRAME_OUTPUT_DIR}/*"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}
